# 1.按照要求，完成VGG16以下处理（每题10分）
# ①数据处理
import tensorflow as tf
from keras.datasets.mnist import load_data
from keras import Model, models, layers, activations, optimizers, losses

# 1)读取mnist数据集
(x_train, y_train), (x_test, y_test) = load_data()
# 2)对数据进行维度转换、归一化等相关预处理
x_train = x_train.reshape(-1, 28, 28, 1) / 255
x_test = x_test.reshape(-1, 28, 28, 1) / 255

print()
# ②设置VGG16模块（类），
# 1)声明一个Sequential包括网络结构中的所有功能层
# 2)根据下图VGG16网络结构构建模型类
# 3)卷积模型取前四组，且初始卷积核个数为16
# 4)每到一个新的卷积组通道数翻倍
# 5)最后三层全链接通道数分别为2048，512，10
# 6)Dropout层失活率设置为0.4
# 7)实现正向传播处理
class VGG16(Model):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conv = models.Sequential([
            layers.Conv2D(filters=16, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=16, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.MaxPooling2D(),

            layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.MaxPooling2D(),

            layers.Conv2D(filters=64, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=64, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=64, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.MaxPooling2D(),

            layers.Conv2D(filters=128, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=128, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=128, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.MaxPooling2D()
        ])

        self.flat = layers.Flatten()

        self.fc = models.Sequential([
            layers.Dense(units=2048, activation=activations.relu),
            layers.Dropout(0.4),
            layers.Dense(units=512, activation=activations.relu),
            layers.Dropout(0.4),
            layers.Dense(units=10, activation=activations.softmax)
        ])

    def call(self, inputs, training=None, mask=None):
        out = self.conv(inputs)
        out = self.flat(out)
        out = self.fc(out)
        return out


# ③	完成模型创建及训练
model = VGG16()
# model.call(x_train)
# model.build(input_shape=(None, 28, 28, 1))
model.compile(optimizer=optimizers.Adam(),
              loss=losses.sparse_categorical_crossentropy,
              metrics='acc')
model.fit(x_train, y_train, epochs=2, batch_size=64)

# ③完成模型创建及训练
# 1)输出模型检验后最终的损失值，准确率
loss, acc = model.evaluate(x_test, y_test)
print('loss:', loss, 'acc:', acc)


